from fastapi import FastAPI, Query, HTTPException, Depends
from typing import List, Optional
from jqdatasdk import *

app = FastAPI(title="JQData API", description="基于聚宽数据的HTTP服务", version="1.0.0")


# 初始化JQData连接
def init_jqdata():
    auth("15810829938", "123QWEqwe")
    return True


# 路由：首页
@app.get("/")
def index():
    return {
        "status": "success",
        "message": "JQData HTTP服务已启动",
        "apis": [
            {"path": "/api/securities", "description": "获取所有证券数据"},
            {
                "path": "/api/stock_info",
                "description": "获取股票信息",
                "params": {"code": "股票代码"},
            },
            {
                "path": "/api/price",
                "description": "获取股票价格",
                "params": {
                    "code": "股票代码",
                    "start_date": "开始日期",
                    "end_date": "结束日期",
                },
            },
        ],
    }


# 路由：获取所有证券数据
@app.get("/api/securities")
def get_securities(
    types: List[str] = Query(None, description="证券类型列表"),
    date: Optional[str] = None,
    _: bool = Depends(init_jqdata),
):
    try:
        # 获取数据
        data = get_all_securities(types=types if types else [], date=date)

        # 将DataFrame转换为JSON
        result = data.reset_index().to_dict(orient="records")

        return {"status": "success", "data": result}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# 路由：获取股票信息
@app.get("/api/stock_info")
def get_stock_info(
    code: str, date: Optional[str] = None, _: bool = Depends(init_jqdata)
):
    try:
        if not code:
            raise HTTPException(status_code=400, detail="请提供股票代码")

        # 获取数据
        data = get_security_info(code)

        # 转换为可序列化的dict
        result = {
            "display_name": data.display_name,
            "name": data.name,
            "start_date": data.start_date.strftime("%Y-%m-%d"),
            "end_date": data.end_date.strftime("%Y-%m-%d") if data.end_date else None,
            "type": data.type,
            "parent": data.parent,
        }

        return {"status": "success", "data": result}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))


# 路由：获取股票价格数据
@app.get("/api/price")
def get_price_data(
    code: str,
    start_date: Optional[str] = None,
    end_date: Optional[str] = None,
    _: bool = Depends(init_jqdata),
):
    try:
        if not code:
            raise HTTPException(status_code=400, detail="请提供股票代码")

        # 获取数据
        data = get_price(
            code,
            start_date=start_date if start_date else None,
            end_date=end_date if end_date else None,
            frequency="daily",
            fields=["open", "close", "high", "low", "volume"],
        )

        # 将DataFrame转换为JSON
        result = data.reset_index().to_dict(orient="records")

        return {"status": "success", "data": result}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
